import math
import numpy as np
import gurobipy as grb

class Routing():
    """
    Desc: 先限流再分流
    """
    _w_routing = {}
    
    def __init__(self, topology_pods):
        """
        Desc: 做routing前，除了需要加载基类的配置读取，配置其他参数
              还需要添加前面的拓扑工程计算出来的pods间连接数
        Inputs:
            - topology_pods(2-dimension list): Number of s_i egress links connected to igress links of s_j.
        """
        self._topology_pods = topology_pods


    def routing(self, Sij, virtual_r, a_link_bandwidth):
        new_Sij = self.scale_Sij(Sij, virtual_r, a_link_bandwidth)
        threshold = new_Sij
        return threshold


    def scale_Sij(self, Sij, virtual_r, a_link_bandwidth):
        record_shape = Sij.shape

        capacities = self._topology_pods * a_link_bandwidth

        m = grb.Model('scale_Sij')
        m.Params.OutputFlag = 0
        alpha = m.addVar(vtype=grb.GRB.CONTINUOUS, lb=0, name='alpha')
        
        m.addConstrs(
            alpha * (Sij[i][j] + virtual_r[i] * virtual_r[j]) <= capacities[i][j]
            for i in range(record_shape[0])
            for j in range(record_shape[1])
            if i != j
        )

        m.setObjective(alpha, grb.GRB.MAXIMIZE)
        m.optimize()
        if m.status == grb.GRB.Status.OPTIMAL:
            alpha_opt = m.objVal
        else:
            print('scale_Sij No solution')
            exit()

        new_Sij = Sij
        for i in range(record_shape[0]):
            for j in range(record_shape[1]):
                if i != j:
                    new_Sij[i][j] = capacities[i][j] / alpha_opt - virtual_r[i] * virtual_r[j]

        return new_Sij * alpha_opt


    def find_threshold(self, traffic, Sij, virtual_r):

        record_shape = Sij.shape

        m = grb.Model('find_threshold')
        m.Params.OutputFlag = 0

        beta = m.addVar(vtype=grb.GRB.CONTINUOUS, name='beta')

        R = virtual_r.sum()

        inter_u = m.addVars(record_shape[0], record_shape[1], lb=0, vtype=grb.GRB.CONTINUOUS, name='u')
      
        m.addConstrs(
            inter_u[i, j] >= traffic[i][j] - beta * Sij[i][j]
            for i in range(record_shape[0])
            for j in range(record_shape[1])
            if i != j
        )

        m.addConstrs(
            grb.quicksum(
                inter_u[i, j] for j in range(record_shape[1]) if i != j
            ) <= beta * virtual_r[i] * (R - virtual_r[i])
            for i in range(record_shape[0])
        )

        m.addConstrs(
            grb.quicksum(
                inter_u[i, j] for i in range(record_shape[1]) if i != j
            ) <= beta * virtual_r[j] * (R - virtual_r[j])
            for j in range(record_shape[1])
        )

        m.setObjective(beta, grb.GRB.MINIMIZE)
        # m.write('find_threshold.lp')
        m.optimize()
        if m.status == grb.GRB.Status.OPTIMAL:
            beta_opt = m.objVal
        else:
            print('find_threshold No solution')
            exit()

        threshold = beta_opt * Sij
        return threshold
